import numpy as np
import torch
import time
import importlib
import argparse
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  #CPU MODEL ONLY



def main(net_name, num_classes, input_shape, test_cnt=10 ):
    module = importlib.import_module("nets.%s.model" % net_name)
    model = module.get_seg_model(num_classes=num_classes, aux=True)
    model = model.to("cpu")
    model.eval()

    img = np.ones((1, input_shape[2], input_shape[1], input_shape[0]), dtype=np.float32)
    input_data = torch.from_numpy(img).to("cpu")
    input_data = input_data.float()

    for t in range(10):
        start_t = time.time()
        for i in range(test_cnt):
            output = model(input_data)
            if isinstance(output, (list, tuple)):
                output_shape = ""
                for _ in output:
                    output_shape += (str(_.shape)+" ")
            else:
                output_shape = str(output.shape)
        end_t = time.time()
        print("%s model run [%d %d %d] use time %.4f ms output_shape:%s" % (
        net_name, input_shape[0], input_shape[1], input_shape[2],
        1000.0 * (end_t - start_t) / test_cnt, output_shape))

if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--net_name", required=True)
    parser.add_argument('--num_classes', type=int, default=1)
    parser.add_argument('--input_shape', type=int, nargs="+", default=(224, 224, 3))
    args = parser.parse_args()
    main(args.net_name, args.num_classes, args.input_shape)
